import os
import sys
from loguru import logger
from starlette.requests import Request
from starlette.responses import Response
from datetime import datetime, timedelta, timezone
from starlette.middleware.base import BaseHTTPMiddleware

# 创建日志文件夹及其子文件夹
log_directory = "logs"
os.makedirs(os.path.join(log_directory, "error"), exist_ok=True)
os.makedirs(os.path.join(log_directory, "info"), exist_ok=True)


class LoggerMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next) -> Response:
        method = request.method
        path = request.scope["path"]
        http_type = request.scope["type"]
        china_tz = timezone(timedelta(hours=8), "China")
        current_time_china = datetime.now(china_tz)
        formatted_time = current_time_china.strftime("%Y-%m-%d")

        # 处理请求
        response = await call_next(request)

        code = response.status_code

        # 此处移出日志管理器，避免后续的日志记录混淆
        logger.remove()

        # 根据状态码选择日志文件夹
        log_folder = "error" if code >= 400 else "info"
        log_file = os.path.join(log_directory, log_folder, f"{formatted_time}.log")

        logger.add(log_file, rotation="1 week", enqueue=True)  # 添加日志文件处理器

        # 记录请求信息
        if code >= 400:
            logger.error(
                f"<red>{method} - {code}</red> - <bold><i>{path} - {http_type}</i></bold> - "
                f"<blue>{current_time_china}</blue> - <green>请求失败</green>"
            )
        else:
            logger.info(
                f"<green>{method} - {code}</green> - <bold><i>{path} - {http_type}</i></bold> - "
                f"<blue>{current_time_china}</blue>"
            )

        return response
